package com.katze.ssh.sftp;

import com.katze.ssh.IChannel;
import com.katze.ssh.SSHException;
import com.katze.ssh.SshSessionEntry;
import org.apache.sshd.common.SshException;
import org.apache.sshd.common.channel.exception.SshChannelOpenException;
import org.apache.sshd.sftp.client.SftpClient;
import org.apache.sshd.sftp.client.SftpErrorDataHandler;
import org.apache.sshd.sftp.client.SftpVersionSelector;
import org.apache.sshd.sftp.client.impl.DefaultSftpClientFactory;
import org.apache.sshd.sftp.common.SftpConstants;
import org.apache.sshd.sftp.common.SftpException;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.Consumer;
import java.util.stream.Stream;
import java.util.zip.ZipOutputStream;

public class SftpChannel implements IChannel {
    public static ConcurrentMap<String, SftpDownloadTask> PROGRESS = new ConcurrentHashMap<>();
    private final SshSessionEntry session;
    private String homePath;
    private SftpClient sftpClient;

    public SftpChannel(SshSessionEntry session) {
        this.session = session;
    }

    @Override
    public void open() throws Exception{
        try {
            this.sftpClient = new DefaultSftpClientFactory().createSftpClient(session.getClientSession(), SftpVersionSelector.CURRENT, SftpErrorDataHandler.EMPTY);
            //获取当前目录的绝对路径。这里的 "." 表示当前目录。
            this.homePath = this.sftpClient.canonicalPath(".");
        } catch (SshException e) {
            if (e.getCause() instanceof TimeoutException) {
                throw new SSHException("会话连接超时");
            } if (e.getCause() instanceof SshChannelOpenException) {
                session.getClientSession().close();
                throw new SSHException("会话连接异常，请重试");
            } else {
                throw e;
            }
        }
    }

    public String getHomePath() {
        return homePath;
    }

    public SftpClient.Attributes exists(String path) {
        try {
            return sftpClient.stat(path);
        } catch (Exception e) {
            if (e instanceof SftpException se) {
                if (SftpConstants.getStatusName(se.getStatus()).equals("SSH_FX_NO_SUCH_FILE")) {
                    return null;
                }
            }
            throw new RuntimeException(e);
        }
    }

    public List<FileEntry> home() {
        List<FileEntry> list = new ArrayList<>();
        try {
            String[] items = this.homePath.split("/");
            for (int i = 0; i < items.length; i++) {
                if ( i == 0) {
                    SftpClient.Attributes attributes = sftpClient.stat("/");
                    FileEntry entry = new FileEntry("/", "/", attributes);
                    list.add(entry);
                } else {
                    FileEntry pro = list.get( i - 1);
                    String path = pro.getFullPath();
                    SftpClient.Attributes attributes = sftpClient.stat(path.equals("/")? path + items[i]: path + "/" + items[i]);
                    FileEntry entry = new FileEntry(path, items[i], attributes);
                    list.add(entry);
                }
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return list;
    }

    public void ls(String dir, BiConsumer<String, Stream<SftpClient.DirEntry>> action) throws Exception {
        dir = toLinuxPath(dir);
        action.accept(dir, this.dirFilter(sftpClient.readEntries(dir)));
    }

    /**
     * 新建文件夹
     * @param dir 路径
     */
    public String mkdir(String dir) throws Exception {
        dir = toLinuxPath(dir);
        SftpClient.Attributes attrs = exists(dir);
        if (attrs != null) {
            int i = 1;
            String newDir = dir + "/新建文件夹";
            while (true) {
                attrs = exists(newDir);
                if (attrs != null && attrs.isDirectory()) {
                    newDir = dir + "/新建文件夹(" + i + ")";
                    i++;
                } else {
                    break;
                }
            }
            sftpClient.mkdir(newDir);
            return newDir;
        } else {
            throw new SSHException(dir + "目录不存在");
        }
    }

    public String mkdirs(String dir, String name) throws Exception {
        dir = toLinuxPath(dir);
        SftpClient.Attributes attrs = exists(dir);
        if (attrs != null) {
            name = toLinuxPath(name);
            String[] names = name.split("/");
            for (String item : names) {
                if (item.isEmpty()) continue;
                attrs = exists(dir += "/" + item);
                if (attrs == null || !attrs.isDirectory()) {
                    sftpClient.mkdir(dir);
                }
            }
            return dir;
        } else {
            throw new SSHException(dir + "目录不存在");
        }
    }

    /**
     * 移动
     * @param source    起始路径
     * @param target    目标路径
     * @param filenames 文件名
     */
    public List<FileEntry> mv(String source, String target, String... filenames) throws Exception {
        target = toLinuxPath(target);
        List<String> existingFilenames;
        try (SftpClient.CloseableHandle handle = sftpClient.openDir(target)){
            existingFilenames = sftpClient.readDir(handle).stream().map(SftpClient.DirEntry::getFilename).toList();
        }

        Map<String, String[]> map = new HashMap<>();
        source = toLinuxPath(source);
        for (String filename : filenames) {
            String name = filename;
            boolean exist;
            int c = 1;
            do {
                exist = existingFilenames.contains(name);
                if (exist) {
                    int index = filename.lastIndexOf(".");
                    String txt = c > 1? " - 副本(" + c + ")":" - 副本";
                    if (index == -1 || index == 0) {
                        name = filename + txt;
                    } else {
                        name = filename.substring(0, index) + txt + filename.substring(index);
                    }
                    c ++;
                }
            } while (exist);
            map.put(source + "/" + filename, new String[]{target, name});
        }
        List<FileEntry> list = new ArrayList<>();
        for (Map.Entry<String, String[]> entry : map.entrySet()) {
            String[] vs = entry.getValue();
            String fullPath = vs[0] + "/" + vs[1];
            sftpClient.rename(entry.getKey(), fullPath);
            SftpClient.Attributes attributes = sftpClient.stat(fullPath);
            list.add(new FileEntry(vs[0], vs[1], attributes));
        }
        return list;
    }

    /**
     * 单文件/目录移动
     * @param source 起始路径
     * @param sourceName 文件名
     * @param target 目标路径
     * @param targetName 目标文件名
     */
    public void mv(String source, String sourceName, String target, String targetName) throws Exception {
        source = toLinuxPath(source) + "/" + sourceName;
        target = toLinuxPath(target) + "/" + targetName;

        if (sftpClient.getVersion() <= 3) {
            sftpClient.remove(target);
        }
        sftpClient.rename(source , target);
    }

    /**
     * 复制
     * @param source    起始路径
     * @param target    目标路径
     * @param filenames 文件名
     */
    public List<FileEntry> cp(String source, String target, String... filenames) throws Exception {
        target = toLinuxPath(target);
        List<String> existingFilenames;
        try (SftpClient.CloseableHandle handle = sftpClient.openDir(target)){
            existingFilenames = sftpClient.readDir(handle).stream().map(SftpClient.DirEntry::getFilename).toList();
        }

        Map<String, String[]> map = new HashMap<>();
        source = toLinuxPath(source);
        for (String filename : filenames) {
            String name = filename;
            boolean exist;
            int c = 1;
            do {
                exist = existingFilenames.contains(name);
                if (exist) {
                    int index = filename.lastIndexOf(".");
                    String txt = c > 1? " - 副本(" + c + ")":" - 副本";
                    if (index == -1 || index == 0) {
                        name = filename + txt;
                    } else {
                        name = filename.substring(0, index) + txt + filename.substring(index);
                    }
                    c ++;
                }
            } while (exist);
            map.put(source + "/" + filename, new String[]{target, name});
        }
        List<FileEntry> list = new ArrayList<>();
        for (Map.Entry<String, String[]> entry : map.entrySet()) {
            String[] vs = entry.getValue();
            String fullPath = vs[0] + "/" + vs[1];
            walkCp(entry.getKey(), fullPath, sftpClient.stat(entry.getKey()));
            SftpClient.Attributes attributes = sftpClient.stat(fullPath);
            list.add(new FileEntry(vs[0], vs[1], attributes));
        }
        return list;
    }

    private void walkCp(String source, String target,SftpClient.Attributes attributes) throws Exception {
        if (attributes.isDirectory()) {
            try (SftpClient.CloseableHandle handle = sftpClient.openDir(source)){
                sftpClient.mkdir(target);
                Iterator<SftpClient.DirEntry> iterator = this.dirFilter(sftpClient.readDir(handle)).iterator();
                while (iterator.hasNext()) {
                    SftpClient.DirEntry entry = iterator.next();
                    walkCp(source + "/" + entry.getFilename(), target + "/" + entry.getFilename(), entry.getAttributes());
                }
            }
        } else {
            try (SftpClient.CloseableHandle sourceHandle = sftpClient.open(source, EnumSet.of(SftpClient.OpenMode.Read));
                 SftpClient.CloseableHandle destHandle = sftpClient.open(target, EnumSet.of(SftpClient.OpenMode.Write, SftpClient.OpenMode.Create, SftpClient.OpenMode.Truncate))) {
                byte[] buffer = new byte[10240];
                int offset = 0;
                for (int length; (length = sftpClient.read(sourceHandle, offset, buffer)) != -1 ; offset += length) {
                    sftpClient.write(destHandle, offset, buffer, 0, length);
                }
            }
        }
    }

    /**
     * 文件重命名
     * @param dir 路径
     * @param oldName 旧文件名
     * @param newName 新文件名
     */
    public SftpClient.Attributes rename(String dir, String oldName, String newName) throws Exception {
        dir = toLinuxPath(dir);
        String oldPath = dir + "/" + oldName;
        SftpClient.Attributes oldAttrs = exists(oldPath);
        if (oldAttrs == null) {
            throw new SSHException(oldPath + "文件不存在");
        }
        String newPath = dir + "/" + newName;
        SftpClient.Attributes newAttrs = exists(newPath);
        if (newAttrs != null) {
            throw new SSHException("目标文件名已存在");
        }
        sftpClient.rename(oldPath, newPath);
        return oldAttrs;
    }

    /**
     * 同目录下删除多个文件/目录
     * @param dir 文件路径
     * @param names 文件名称
     */
    public Consumer<BiConsumer<Integer, Integer>> rmf(String dir, String... names) throws Exception {
        dir = toLinuxPath(dir);
        AtomicInteger dirCount = new AtomicInteger(0);
        AtomicInteger docCount = new AtomicInteger(0);
        for (String name : names) {
            if (name.isBlank()) continue;
            String file = dir + "/" + name;
            walkRmf(file, sftpClient.stat(file), isDir -> {
                if (isDir) {
                    dirCount.getAndIncrement();
                } else {
                    docCount.getAndIncrement();
                }
            });
        }
        return cm -> cm.accept(dirCount.get(), docCount.get());
    }

    private void walkRmf(String file, SftpClient.Attributes attributes, Consumer<Boolean> cm) throws Exception {
        if (attributes.isDirectory()) {
            try (SftpClient.CloseableHandle handle = sftpClient.openDir(file)){
                Iterator<SftpClient.DirEntry> iterator = this.dirFilter(sftpClient.readDir(handle)).iterator();
                while (iterator.hasNext()) {
                    SftpClient.DirEntry entry = iterator.next();
                    walkRmf(file + "/" + entry.getFilename(), entry.getAttributes(), cm);
                }
                sftpClient.rmdir(file);
                cm.accept(true);
            }
        } else {
            sftpClient.remove(file);
            cm.accept(false);
        }
    }

    /**
     * 删除单文件
     * @param dir 目录
     * @param names 文件名
     */
    public void rm (String dir, String... names) throws Exception {
        dir = toLinuxPath(dir);
        for (String name : names) {
            sftpClient.remove(dir + "/" + name);
        }
    }

    /**
     * 下载文件/目录
     * @param remote 远程路径
     * @param local 本地路径
     */
    public void download(String remote, String local) throws Exception {
        remote = toLinuxPath(remote);
        Path dir = Paths.get(local);
        if (!Files.exists(dir)) {
            Files.createDirectories(dir);
        }
        walkRemote(remote, sftpClient.stat(remote), dir);
    }

    private void walkRemote(String remote, SftpClient.Attributes attributes, Path local) throws Exception{
        if (attributes.isDirectory()) {
            List<SftpClient.DirEntry> list = this.dirFilter(sftpClient.readEntries(remote)).toList();
            for (SftpClient.DirEntry entry : list) {
                if (entry.getAttributes().isDirectory()) {
                    walkRemote(remote + "/" + entry.getFilename(), entry.getAttributes(), local.resolve(entry.getFilename()));
                }else {
                    walkRemote(remote + "/" + entry.getFilename(), entry.getAttributes(), local);
                }
            }
        } else {
            if (!Files.exists(local)) {
                Files.createDirectories(local);
            }
            try (InputStream in = sftpClient.read(remote)){
                String filename = remote.substring(remote.lastIndexOf(47) + 1);
                Files.copy(in, local.resolve(filename));
            }
        }
    }

    /**
     *
     * @param id 下载标识
     * @param out 流
     * @param data 下载文件（key:路径- value:文件名）
     */
    public void download(String id, OutputStream out, Map<String,List<String>> data) throws Exception {
        if (data == null || data.isEmpty()) return;
        try (ZipOutputStream stream = new ZipOutputStream(out); SftpDownloadTask task = PROGRESS.computeIfAbsent(id, SftpDownloadTask::new)){
            for (Map.Entry<String, List<String>> entry : data.entrySet()) {
                String source = toLinuxPath(entry.getKey());
                List<String> filenames = entry.getValue();
                for (String filename : filenames) {
                    walkZip(source, filename, sftpClient.stat(source + "/" + filename), task::add);
                }

                //获取当前 SSH 会话中客户端发送数据的最大数据包大小
//                long size = sftpClient.getClientChannel().getLocalWindow().getPacketSize();
                byte[] buffer = new byte[1024 * 1024];
                for (ZipEntity item; (item = task.poll()) != null; task.done()) {
                    stream.putNextEntry(item);
                    if (item.isDirectory()) {
                        item.setOffset((int) item.getSize());
                        stream.closeEntry();
                    } else {
                        // 方式1
                        try (FileChannel channel = sftpClient.openRemoteFileChannel(item.getFilePath(), EnumSet.of(SftpClient.OpenMode.Read))) {
                            ByteBuffer byteBuffer = ByteBuffer.wrap(buffer);
                            // 从 FileChannel 读取数据并写入 ZIP 流
                            for (int size; (size = channel.read(byteBuffer)) != -1; item.setOffset(size)) {
                                stream.write(buffer, 0, size);
                                byteBuffer.clear(); // 清空缓冲区，准备下一次读取
                            }
                        } finally {
                            stream.closeEntry();
                        }
                        // 方式2
                        /*try (SftpClient.CloseableHandle handle = sftpClient.open(item.getFilePath(), EnumSet.of(SftpClient.OpenMode.Read));) {
                            for (int length; (length = sftpClient.read(handle, item.getOffset(), buffer)) != -1;) {
                                // 每次写入大小，会受到带宽影响
                                stream.write(buffer, 0, length);
                                item.setOffset(length);
                            }
                        } finally {
                            stream.closeEntry();
                        }*/
                    }
                }
            }
        } finally {
            PROGRESS.remove(id);
        }
    }

    private void walkZip(String source, String filename, SftpClient.Attributes attributes, Consumer<ZipEntity> con) {
        if (attributes.isDirectory()) {
            try {
                List<SftpClient.DirEntry> list = this.dirFilter(sftpClient.readEntries(source + "/" + filename)).toList();
                for (SftpClient.DirEntry entry : list) {
                    String name = filename + "/" + entry.getFilename();
                    if (entry.getAttributes().isDirectory()) {
                        con.accept(new ZipEntity(source, name, entry.getAttributes()));
                    }
                    walkZip(source, name, entry.getAttributes(), con);
                }
            } catch (Exception e) {
                throw new SSHException(source + "/" + filename + "加载失败", e);
            }
        } else {
            con.accept(new ZipEntity(source, filename, attributes));
        }
    }

    /**
     * 上传文件/目录
     * @param local 本地路径
     * @param remote 远程路径
     */
    public void put(Object local, String remote) throws Exception {
        Path target;
        if (local instanceof Path tg) {
            target = tg;
        } else {
            target = Paths.get(local.toString());
        }
        walkLocal(target, Files.isDirectory(target), toLinuxPath(remote));
    }

    private void walkLocal(Path local, boolean isDir, String remote) throws Exception{
        if (isDir) {
            List<Path> list = Files.list(local).toList();
            for (Path item : list) {
                if (Files.isDirectory(item)) {
                    walkLocal(item, true, remote + "/" + item.getFileName());
                } else {
                    walkLocal(item, false, remote);
                }
            }
        } else {
            if (exists(remote) == null) {
                sftpClient.mkdir(remote);
            }
            sftpClient.put(local, remote + "/" + local.getFileName());
        }
    }

    /**
     * 上传文件
     * @param remote 远程路径
     * @param action 推送文件内容
     */
    public void put(String remote, Consumer<FileChannel> action) throws Exception {
        remote = toLinuxPath(remote);
        try (FileChannel channel = sftpClient.openRemoteFileChannel(remote, EnumSet.of(SftpClient.OpenMode.Write, SftpClient.OpenMode.Create, SftpClient.OpenMode.Append))){
            action.accept(channel);
        }
    }

    private String toLinuxPath(String str) {
        if (str != null && !str.isBlank()) {
            byte[] bytes = str.getBytes(StandardCharsets.UTF_8);
            int i;
            for (i = 0; i < bytes.length; i++) {
                if (bytes[i] == 92) {
                    bytes[i] = 47;
                }
            }
            if (i > 1 && bytes[i - 1] == 47) {
                byte[] temp = new byte[bytes.length - 1];
                System.arraycopy(bytes, 0, temp, 0, temp.length);
                bytes = temp;
            }
            return new String(bytes);
        }
        throw new SSHException("路径不能为空");
    }

    private Stream<SftpClient.DirEntry> dirFilter(Collection<SftpClient.DirEntry> list) {
        return list.stream().filter(item -> !item.getFilename().equals(".") && !item.getFilename().equals(".."));
    }

    @Override
    public void close() throws IOException {
        if (this.sftpClient.isOpen()) {
            this.sftpClient.close();
        }
    }
}
